package com.tjc.loadbalance;

import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.*;

/**
 * @Author: tjc
 * @description
 * @Date Created in 18:05 2022-02-01
 */
public class LB {
    private final int workerNum;
    Set<Worker> workers = null;
    BlockingQueue<Runnable> taskQueue = new ArrayBlockingQueue<>(1000);
    private Runnable toDealTask = null;
    private Phaser phaser = new Phaser();
    private Thread listener;

    private void shutdown() {
        workers.stream().forEach(Worker::shutdown);
        listener.interrupt();
    }

    public LB(int... workParallel) {
        this.workerNum = workParallel.length;
        this.workers = new HashSet<>(workerNum);
        for (int i = 0; i < workParallel.length; i++) {
            workers.add(new Worker(workParallel[i], "worker-" + i));
        }
        listener = new Thread(() -> {
            // cas阻塞而不是轮询
            try {
                while (!Thread.interrupted()) {

                    if (toDealTask == null) {
                        toDealTask = taskQueue.take();
                    }
                    if (assign(toDealTask)) {
                        toDealTask = null;
                    }
                }
            } catch (InterruptedException e) {
                System.out.println("listener is closed.");
            }
        });
        listener.start();
    }

    private boolean assign(Runnable task) {
        Iterator<Worker> itr = workers.iterator();
        while (itr.hasNext()) {
            Worker worker = itr.next();
            try {
                if (worker.accept(task))
                    return true;
            } catch (InterruptedException e) {
                worker.shutdown();
                itr.remove();
                e.printStackTrace();
            }
        }
        return false;
    }

    private void submit(Runnable cmd) {
        // todo submit不能阻塞
        Task task = new Task(cmd);
        try {
            phaser.register();
//            System.out.println(String.format("phaser num:%s level: %s", phaser.getRegisteredParties(), phaser.getPhase()));
            taskQueue.add(task);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public void await() {
        phaser.awaitAdvance(0);
    }

    class Task implements Runnable {
        public static final int READY = 0;
        public static final int RUNNING = 1;
        public static final int FINISH = 2;
        public static final int FAIL = 3;
        int status;
        Runnable firstTask;

        public Task(Runnable firstTask) {
            this.firstTask = firstTask;
        }

        public void beforeExec() {
        }

        public void afterExec() {
            phaser.arrive();
//            System.out.println(String.format("phaser arrive:%s level: %s", phaser.getArrivedParties(), phaser.getPhase()));
        }

        @Override
        public void run() {
            beforeExec();
            runTask();
            afterExec();
        }

        private void runTask() {
            firstTask.run();
        }
    }

    class Worker {
        private String name;
        private int parallel = 1;
        private Set<Thread> threads = new HashSet<>();
        private BlockingQueue<Runnable> cmdQueue = new LinkedBlockingQueue<>();
        private Semaphore lock;
        private Runnable func = () -> {
            try {
                while (!Thread.interrupted()) {
                    Runnable cmd = cmdQueue.take();
                    cmd.run();
                    lock.release();
                }
            } catch (InterruptedException e) {
                System.out.println(Thread.currentThread().getName() + " is closed.");
            }
        };

        public Worker(int parallel, String name) {
            this.name = name;
            this.parallel = parallel;
            this.lock = new Semaphore(parallel);
            for (int i = 0; i < parallel; i++) {
                Thread t = new Thread(func, this.name + "-" + i);
                threads.add(t);
                t.start();
            }
            System.out.println(String.format("worker[%s] started, parallel=%s", name, this.parallel));
        }

        public boolean accept(Runnable task) throws InterruptedException {
            if (lock.tryAcquire()) {
                cmdQueue.put(task);
                return true;
            } else {
                return false;
            }
        }

        public void shutdown() {
            for (Thread t : threads) {
                t.interrupt();
            }
        }

        public int getParallel() {
            return this.parallel;
        }
    }

    public static void main(String[] args) throws InterruptedException {
        ExecutorService pool = Executors.newFixedThreadPool(5);
        long start1 = System.currentTimeMillis();
        for (int i = 0; i < 10; i++) {
            int finalI = i;
            pool.submit(() -> {
                try {
                    Thread.sleep(1000);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
                System.out.print(finalI);
            });
        }
        pool.shutdown();
        while (!pool.isTerminated()) ;
        long end1 = System.currentTimeMillis();
        System.out.println(" hreadpool cost " + (end1 - start1) + "ms");


        //性能要接近jdk的线程池吧2015ms

        LB lb = new LB(1, 1, 1, 2);
        long start = System.currentTimeMillis();
        for (int i = 0; i < 10; i++) {
            int finalI = i;
            lb.submit(() -> {
                try {
                    Thread.sleep(1000);
                } catch (InterruptedException e) {
                    e.printStackTrace();
                }
                System.out.println(Thread.currentThread().getName() + " execute " + finalI);
            });
        }
        lb.await();
        long end = System.currentTimeMillis();
        System.out.println(end - start + "ms");
        lb.shutdown();
    }
}